#!/usr/bin/env python3
# -*- encoding: utf-8 -*-


import numpy as np
from pathlib import Path
import pickle


work_dir = Path(__file__).resolve().parent
save_path = work_dir / 'params.pkl'


def init_params(input_size, hidden_size, output_size, weight_init_std=0.01):
    params = {
        'W1': weight_init_std * np.random.randn(input_size, hidden_size),
        'b1': np.zeros(hidden_size),
        'W2': weight_init_std * np.random.randn(hidden_size, output_size),
        'b2': np.zeros(output_size)
    }
    data = {
        'dim': (input_size, hidden_size, output_size),
        'params': params
    }
    with save_path.open('wb') as f:
        pickle.dump(data, f)
    return params


def load_params(input_size, hidden_size, output_size, weight_init_std=0.01):
    print(work_dir)
    print(save_path)
    if save_path.exists():
        with save_path.open('rb') as f:
            data = pickle.load(f)
        if data['dim'] == (input_size, hidden_size, output_size):
            params = data['params']
        else:
            params = init_params(input_size, hidden_size, output_size)
    else:
        params = init_params(input_size, hidden_size, output_size)
    return params
